#include "tgaimage.h"
#include <fstream>
#include <iostream>
#include <memory>
#include <cstring>
#include <cassert>
#include <cerrno>

namespace raster {

TGAImage::TGAImage()
: w(0), h(0), bpp(0), data(0) {
    header = std::make_shared<TGAHeader>();
    extensionArea = std::make_shared<TGAExtensionArea>();
    developeArea = std::make_shared<TGADevelopArea>();
    footer = std::make_shared<TGAFooter>(); 
}

TGAImage::TGAImage(int width, int height, int bpp)
: w(width), h(height), bpp(bpp), data(width * height * bpp, 0) {
    header = std::make_shared<TGAHeader>();
    extensionArea = std::make_shared<TGAExtensionArea>();
    developeArea = std::make_shared<TGADevelopArea>();
    footer = std::make_shared<TGAFooter>(); 
}

TGAImage::~TGAImage() {
    data.clear();
}

bool TGAImage::readTgaFile(const char* filename) {
    std::ifstream in(filename, std::ios::binary);
    if (!in.is_open()) {
        std::cerr << "Failed to open file: " << filename << " for reading" << std::endl;
        char buf[1024];
#ifdef _WIN32
        strerror_s(buf, sizeof(buf), errno);
#else // POSIX.1-2001
        strerror_r(errno, buf, sizeof(buf));
#endif
        std::cerr << "reason(errno=" << errno << "): " << buf << std::endl;
        return false;
    }

    filePath = filename;

    // don't change the order of reading
    // read header
    if (!readHeader(in)) {
        std::cerr << "An error occurred while reading header of file " << filename << std::endl;
        return false;
    }

    // read data
    if (!readData(in)) {
        std::cerr << "An error occurred while reading data of file " << filename << std::endl;
        return false;
    }

    // read footer
    if (!readFooter(in)) {
        std::cerr << "An error occurred while reading footer of file " << filename << std::endl;
        return false;
    }

    // !!!warning: the following areas need to read after reading footer
    // read develop area
    if (!readDevelopArea(in)) {
        std::cerr << "An error occurred while reading develop area of file " << filename << std::endl;
        return false;
    }

    // read extension area
    if (!readExtensionArea(in)) {
        std::cerr << "An error occurred while reading extension area of file " << filename << std::endl;
        return false;
    }

    // flip image
    if (header->imageDescriptor & 0x20) {
        // bit5: origin to screen order: top to bottom
        flipVertically();
    }
    if (header->imageDescriptor & 0x10) {
        // bit4: pixel to screen order: left to right
        flipHorizontally();
    }
    
    std::cerr << w << "x" << h << "/" << bpp * 8 << "bits" << std::endl;
    in.close();
    return true;
}

bool TGAImage::writeTgaFile(const char* filename, bool vflip, bool rle) {
    std::ofstream out(filename, std::ios::binary | std::ios::out | std::ios::trunc);
    if (!out.is_open()) {
        std::cerr << "Failed to open file: " << filename << " for writing" << std::endl;
        char buf[1024];
#ifdef _WIN32
        strerror_s(buf, sizeof(buf), errno);
#else
        strerror_r(errno, buf, sizeof(buf));
#endif
        std::cerr << "reason(errno=" << errno << "): " << buf << std::endl;
        return false;
    }

    filePath = filename;

    // write header
    if (!writeHeader(out, vflip, rle)) {
        std::cerr << "An error occurred while writing header of file " << filename << std::endl;
        return false;
    }

    // write data
    if (!writeData(out)) {
        std::cerr << "An error occurred while writing data of file " << filename << std::endl;
        return false;
    }

    // write develop area
    if (!writeDevelopArea(out)) {   
        std::cerr << "An error occurred while writing develop area of file " << filename << std::endl;
        return false;
    }

    // write extension area
    if (!writeExtensionArea(out)) {
        std::cerr << "An error occurred while writing extension area of file " << filename << std::endl;
        return false;
    }

    // write footer
    if (!writeFooter(out)) {
        std::cerr << "An error occurred while writing footer of file " << filename << std::endl;
        return false;
    }
    
    out.close();
    return true;
}

TGAColor TGAImage::get(int x, int y) const {
    TGAColor res;
    if (data.empty() || x < 0 || x >= w || y < 0 || y >= h) {
        return res;
    }
    
    memcpy(res.bgra, data.data() + (x + y * w) * bpp, bpp);
    return res;
}

void TGAImage::set(int x, int y, const TGAColor &c) {
    if (data.empty() || x < 0 || x >= w || y < 0 || y >= h) {
        return;
    }
    memcpy(data.data() + (x + y * w) * bpp, c.bgra, bpp);
}

void TGAImage::flipHorizontally() {
    // x + j = w - 1 => x = w - 1 - j
    for (int i = 0; i < h; i++) {
        for (int j = 0; j < w / 2; j++) {
            for (int k = 0; k < bpp; k++) {
                std::swap(data[i * w * bpp + j * bpp + k], data[i * w * bpp + (w - j - 1) * bpp + k]);
            }
        }
    }
}

void TGAImage::flipVertically() {
    // y + i = h - 1 => y = h - 1 - i
    for (int i = 0; i < h / 2; i++) {
        for (int j = 0; j < w; j++) {
            for (int k = 0; k < bpp; k++) {
                std::swap(data[i * w * bpp + j * bpp + k], data[(h - i - 1) * w * bpp + j * bpp + k]);
            }
        }
    }
}

bool TGAImage::readHeader(std::ifstream &in) {
    if (!in.is_open()) {
        std::cerr << "Please open file " << filePath << " for reading first" << std::endl;
        return false;
    }

    in.read(reinterpret_cast<char*>(header.get()), sizeof(TGAHeader));
    if (!in.good()) {
        return false;
    }

    w = header->width;
    h = header->height;
    bpp = header->bitsPerPixel >> 3;
    data.resize(w * h * bpp);
    if (w <= 0 || h <= 0 || (bpp != GRAYSCALE && bpp != RGB && bpp != RGBA)) {
        std::cerr << "Invalid image size or bit depth: " << w << "x" << h << "x" << bpp << std::endl;
        return false;
    }

    return true;
}

bool TGAImage::writeHeader(std::ofstream &out, bool vflip, bool rle) {
    if (!out.is_open()) {
        std::cerr << "Please open file " << filePath << " for writing first" << std::endl;
        return false;
    }

    // not support color map data for now

    bool hflip = false;
    header->idLength = 0;
    header->colorMapType = 0;
    header->imageType = (bpp == GRAYSCALE) ? (rle ? 11 : 3) : (rle ? 2 : 10);
    header->firstEntryIndex = 0;
    header->colorMapLength = 0;
    header->colorMapEntrySize = 0;
    header->xOrigin = 0;
    header->yOrigin = 0;
    header->width = w;
    header->height = h;
    header->bitsPerPixel = bpp << 3;
    header->imageDescriptor = (vflip ? 0x20 : 0) | (hflip ? 0x10 : 0);

    out.write(reinterpret_cast<char*>(header.get()), sizeof(TGAHeader));
    if (!out.good()) {
        std::cerr << "Can't write header" << std::endl;
        return false;
    }

    return true;
}

bool TGAImage::readData(std::ifstream &in) {
    assert(data.size() == w * h * bpp); // check data size

    if (!in.is_open()) {
        std::cerr << "Please open file " << filePath << " for reading first" << std::endl;
        return false;
    }

    // read image id
    std::vector<uint8_t> imageId(header->idLength, 0);
    if (header->idLength > 0) {
        in.read(reinterpret_cast<char*>(imageId.data()), header->idLength);
        if (!in.good()) {
            std::cerr << "An error occurred while reading image id" << std::endl;
            return false;
        }
    }

    // read data: Image Data, or Color Map Data
    switch (header->imageType) {
        // 0~127: Truevision Co. ; 128~255: deverloper
        case 0: // No Image Data Included
            data.clear();
            break;

        case 1: // Uncompressed, Color-mapped Image
            // required: header.colorMapType == 1

            // TODO: to be completed
            break;

        case 2: // Uncompressed, True-color Image
            // required: header.colorMapType == 0
        case 3: // Uncompressed, Black-and-white Image
            // required: header.colorMapType == 0

            // image data
            in.read(reinterpret_cast<char*>(data.data()), data.size());
            if (!in.good()) {
                std::cerr << "An error occurred while reading image data of file " <<  std::endl;
                return false;
            }
            break;

        case 9: // Run-length encoded, Color-mapped Image
            // required: header.colorMapType == 1

            // TODO: to be completed
            break;

        case 10: // Run-length encoded, True-color Image
            // required: header.colorMapType == 0
        case 11: // Run-length encoded, Black-and-white Image
            // required: header.colorMapType == 0

            // colormap data
            if (!loadRleData(in)) {
                std::cerr << "An error occurred while reading RLE image data of file " << std::endl;
                return false;
            }
            break;
        
        default:
            std::cerr << "Unsupported TGA image type: " << header->imageType << std::endl;
            return false;
    }

    return true;
}

bool TGAImage::writeData(std::ofstream &out, bool rle) {
    if (!out.is_open()) {
        std::cerr << "Please open file " << filePath << " for writing first" << std::endl;
        return false;
    }

    // !!!warning: not support color map data for now
    // 1. ignore image id

    // 2. write image data, or color map data according to image type
    switch (header->imageType) {
        case 0: // No Image Data Included
            break;

        case 1: // Uncompressed, Color-mapped Image
            // TODO: to be completed
            break;

        case 2: // Uncompressed, True-color Image
        case 3: // Uncompressed, Black-and-white Image

            out.write(reinterpret_cast<char*>(data.data()), data.size());
            if (!out.good()) {
                    std::cerr << "An error occurred while writing image data" << std::endl;
                return false;
            }
            break;
            
        case 9: // Run-length encoded, Color-mapped Image
            // TODO: to be completed
            break;

        case 10: // Run-length encoded, True-color Image
        case 11: // Run-length encoded, Black-and-white Image
            unloadRleData(out);
            break;

        default:
            std::cerr << "Unsupported TGA image type: " << header->imageType << std::endl;
            break;
    }

    return true;
}

bool TGAImage::readDevelopArea(std::ifstream &in) {
    if (!in.is_open()) {
        std::cerr << "Please open file " << filePath << " for reading first" << std::endl;
        return false;
    }

    if (footer->developerDirectoryOffset == 0) { // no developer directory
        return true;
    }

    //  read developer directory
    #if 0 // TODO: read developer directory
    uint16_t numTags = 0;
    in.seekg(footer->developerDirectoryOffset, std::ios::beg);
    in.read(reinterpret_cast<char*>(&numTags), sizeof(numTags));
    for (int i = 0; i < numTags; i++) {

    }

    if (!in.good()) {
        std::cerr << "An error occurred while reading develop area" << std::endl;
        return false;
    }

    //  read developer fields
    in.seekg(footer->developerDirectoryOffset + developeArea->developeFields.size(), std::ios::beg);
    in.read(reinterpret_cast<char*>(developeArea->developeFields.data()), developeArea->developeFields.size());
    if (!in.good()) {
        std::cerr << "An error occurred while reading developer fields" << std::endl;
        return false;
    }
    #endif

    return true;
}

bool TGAImage::writeDevelopArea(std::ofstream &out) {
    if (!out.is_open()) {
        std::cerr << "Please open file " << filePath << " for writing first" << std::endl;
        return false;
    }

    // no develop area for now
    footer->developerDirectoryOffset = 0;
    return true;
}

bool TGAImage::readExtensionArea(std::ifstream &in) {
    if (!in.is_open()) {
        std::cerr << "Please open file " << filePath << " for reading first" << std::endl;
        return false;
    }

    if (footer->extensionOffset == 0) { // no extension area
        return true;
    }

    //  1. read fields with fixed position and size
    in.seekg(footer->extensionOffset, std::ios::beg);
    in.read(reinterpret_cast<char*>(extensionArea.get()), kTgaExtensionAreaFixedSize);
    if (!in.good()) {
        std::cerr << "An error occurred while reading extension area" << std::endl;
        return false;
    }

    // 2. read color correction table
    if (extensionArea->colorCorrectionOffset != 0) {
        in.seekg(extensionArea->colorCorrectionOffset, std::ios::beg);
        in.read(reinterpret_cast<char*>(extensionArea->colorCorrectionTable.data()), extensionArea->colorCorrectionTable.size());
        if (!in.good()) {
            std::cerr << "An error occurred while reading color correction table" << std::endl;
            return false;
        }
    }

    //  3. read scan line table
    if (extensionArea->scanLineOffset != 0) {
        size_t nScanLines = h;
        in.seekg(extensionArea->scanLineOffset, std::ios::beg);
        extensionArea->scanLineTable.resize(nScanLines * sizeof(uint32_t));
        in.read(reinterpret_cast<char*>(extensionArea->scanLineTable.data()), nScanLines * sizeof(uint32_t));
        if (!in.good()) {
            std::cerr << "An error occurred while reading scan line table" << std::endl;
            return false;
        }
    }

    //  4. read postage stamp image
    if (extensionArea->pastageStampOffset != 0) {
        int xSize = 0, ySize = 0;
        in.seekg(extensionArea->pastageStampOffset, std::ios::beg);
        in.read(reinterpret_cast<char*>(&xSize), sizeof(xSize));
        in.read(reinterpret_cast<char*>(&ySize), sizeof(ySize));
        extensionArea->postageStampImage.resize(xSize * ySize * bpp);
        in.read(reinterpret_cast<char*>(extensionArea->postageStampImage.data()), extensionArea->postageStampImage.size());
        if (!in.good()) {
            std::cerr << "An error occurred while reading postage stamp image" << std::endl;
            return false;
        }
    }

    return true;
}

bool TGAImage::writeExtensionArea(std::ofstream &out) {
    if (!out.is_open()) {
        std::cerr << "Please open file " << filePath << " for writing" << std::endl;
        return false;
    }

    // no extension area for now
    footer->extensionOffset = 0;

    return true;
}

bool TGAImage::readFooter(std::ifstream &in) {
    if (!in.is_open()) {
        std::cerr << "Please open file " << filePath << " for reading first" << std::endl;
        return false;
    }

    in.seekg(-sizeof(TGAFooter), std::ios::end);
    in.read(reinterpret_cast<char*>(footer.get()), sizeof(TGAFooter));
    if (!in.good()) {
        std::cerr << "An error occurred while reading footer" << std::endl;
        return false;
    }

    // check TGA version
    if (footer->signature[16] != '.') {
        std::cerr << "Invalid signature: wrong TGA file" << footer->signature << std::endl;
        return false;
    }
    if (std::strcmp(footer->signature, kFooterSignatureV2) == 0) {
        tgaVerion = 2; // NEW TGA
    }
    else {
        tgaVerion = 1; // OLD TGA
    }
    return true;
}

bool TGAImage::writeFooter(std::ofstream &out)
{
    // if need,
    // 1. set developer directory offset
    // 2. set extension area offset

    // 3. set footer signature
    tgaVerion = 2;
    size_t buflen = sizeof(footer->signature);
    if (tgaVerion == 2) {
        strncpy(footer->signature, kFooterSignatureV2, buflen - 1);
    }
    else {
        strncpy(footer->signature, kFooterSignatureV1, buflen - 1);
    }
    footer->signature[buflen - 1] = '\0';

    out.write(reinterpret_cast<char*>(footer.get()), sizeof(TGAFooter));
    if (!out.good()) {
        std::cerr << "An error occurred while writing footer" << std::endl;
        return false;
    }
    return true;
}

bool TGAImage::loadRleData(std::ifstream& in) {
    assert(data.size() == w * h * bpp); // check data size

    size_t pixelCount = w * h;
    size_t currentPixel = 0;

    while (currentPixel < pixelCount) {
        uint8_t chunkHeader = 0;
        if (!in.read(reinterpret_cast<char*>(&chunkHeader), 1)) {
            return false;
        }
        
        if (chunkHeader < 128) { // Raw chunk, b7 = 0
            chunkHeader++; // repetition count + 1
            for (int i = 0; i < chunkHeader; i++) {
                if (!in.read(reinterpret_cast<char*>(data.data() + currentPixel * bpp), bpp)) { // read pixel value
                    return false;
                }
                currentPixel++;
            }
        } else { // RLE chunk, b7 = 1
            chunkHeader -= 127; // repetition count + 1
            std::vector<uint8_t> buffer(bpp);
            if (!in.read(reinterpret_cast<char*>(buffer.data()), bpp)) { // read pixel value
                return false;
            }
            
            for (int i = 0; i < chunkHeader; i++) {
                std::copy(buffer.begin(), buffer.end(), data.begin() + currentPixel * bpp);
                currentPixel++;
            }
        }
    }
    return true;
}

bool TGAImage::unloadRleData(std::ofstream &out) {
    assert(data.size() == w * h * bpp); // check data size

    if (!out.is_open()) {
        std::cerr << "Please open file " << filePath << " for writing first" << std::endl;
        return false;
    }
    const int maxChunkLength = 128;
    const size_t pixelCount = w * h;
    size_t currentPixel = 0;

    while (currentPixel < pixelCount) {
        size_t chunkStart = currentPixel;
        size_t runLength = 1;
        bool isRLE = false;

        // Look ahead to determine if we should use RLE or raw chunk
        while (runLength < maxChunkLength && (chunkStart + runLength) < pixelCount) {
            bool pixelsEqual = true;
            // 1 pixel occupies bpp bytes, including [a,r,g,b] value
            for (int b = 0; b < bpp; b++) {
                if (data[(chunkStart * bpp) + b] != data[((chunkStart + runLength) * bpp) + b]) {
                    pixelsEqual = false;
                    break;
                }
            }
            if (!pixelsEqual) break;
            runLength++;
        }

        // If we found a run of 2 or more identical pixels, use RLE
        if (runLength > 1) {
            isRLE = true;
        } else {
            // Find length of raw chunk (up to maxChunkLength pixels)
            runLength = 1;
            while (runLength < maxChunkLength && (chunkStart + runLength) < pixelCount) {
                bool nextPixelStartsRun = false;
                if ((chunkStart + runLength + 1) < pixelCount) {
                    nextPixelStartsRun = true;
                    for (int b = 0; b < bpp; b++) {
                        if (data[((chunkStart + runLength) * bpp) + b] != 
                            data[((chunkStart + runLength + 1) * bpp) + b]) {
                            nextPixelStartsRun = false;
                            break;
                        }
                    }
                }
                if (nextPixelStartsRun) break;
                runLength++;
            }
        }

        // Write chunk header
        uint8_t chunkHeader = isRLE ? (uint8_t)(runLength + 127) : (uint8_t)(runLength - 1);
        out.write(reinterpret_cast<char*>(&chunkHeader), 1);

        // Write pixel data
        if (isRLE) {
            // For RLE chunk, write the pixel value once
            out.write(reinterpret_cast<char*>(&data[chunkStart * bpp]), bpp);
        } else {
            // For raw chunk, write all pixel values
            out.write(reinterpret_cast<char*>(&data[chunkStart * bpp]), runLength * bpp);
        }

        if (!out.good()) {
            std::cerr << "Failed to write RLE data" << std::endl;
            return false;
        }

        currentPixel += runLength;
    }

    return true;
}

} // namespace raster
